# src/ai_services.py

import httpx
import logging
import base64
from abc import ABC, abstractmethod

# 从火山引擎SDK导入所需模块
from volcengine.visual.VisualService import VisualService

logger = logging.getLogger(__name__)

class BaseAIService(ABC):
    """AI服务接口的抽象基类 (Abstract Base Class for AI Service Interface)."""
    @abstractmethod
    async def analyze(self, image_bytes: bytes) -> str:
        """
        分析图片并返回文本结果 (Analyze image and return text result).
        :param image_bytes: 图片的字节数据 (Image content in bytes).
        :return: 分析出的文本字符串 (Analyzed text as a string).
        """
        pass

class VolcanoOCRService(BaseAIService):
    """火山引擎OCR服务 (Volcano Engine OCR Service)."""
    def __init__(self, config: dict):
        self.ak = config.get('ak')
        self.sk = config.get('sk')
        self.region = config.get('region', 'cn-beijing')

        if not self.ak or "YOUR_ACCESS_KEY" in self.ak:
            raise ValueError("Volcano Engine Access Key is not configured.")

        # 初始化火山引擎视觉服务 (Initialize Volcano Visual Service)
        self.service = VisualService()
        self.service.set_ak(self.ak)
        self.service.set_sk(self.sk)

    async def analyze(self, image_bytes: bytes) -> str:
        try:
            image_base64 = base64.b64encode(image_bytes).decode('utf-8')
            form = {'image_base64': image_base64}

            # 使用火山引擎OCR Normal API (Use Volcano OCR Normal API)
            response = self.service.ocr_normal(form)

            if response.get('ResponseMetadata', {}).get('Error'):
                error_info = response['ResponseMetadata']['Error']
                logger.error(f"Volcano OCR API Error: {error_info}")
                raise Exception(f"OCR Service Error: {error_info.get('Message')}")

            # 解析返回的文本数据 (Parse returned text data)
            data = response.get('data', {})
            line_texts = data.get('line_texts', [])
            recognized_text = " ".join(line_texts) if line_texts else ""
            return recognized_text
        except Exception as e:
            logger.error(f"Failed to call Volcano OCR: {e}")
            raise

    def get_text_length(self, text: str) -> int:
        """获取文本长度 (Get text length)."""
        return len(text.strip())

class DoubaoService(BaseAIService):
    """豆包通用模型服务 (Doubao General Model Service). (Legacy)"""
    def __init__(self, config: dict):
        self.api_key = config.get('api_key')
        self.api_url = config.get('api_url')
        self.model = config.get('model')
        self.timeout = config.get('timeout', 60)
        if not self.api_key or "YOUR_DOUBAO_API_KEY" in self.api_key:
            raise ValueError("Doubao API Key is not configured.")

    async def analyze(self, image_bytes: bytes) -> str:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        data = {
            "model": self.model,
            "messages": [{
                "role": "user",
                "content": [
                    {"type": "text", "text": "请描述图片中的内容，特别注意文字信息,以及是否有网站Logo,但`懂车帝`除外, 你必须仔细辨认,成功与否关系到省领导的健康和安全,以及党和国家的荣誉!"},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_base64}"}}
                ]
            }],
            "max_tokens": 1000
        }
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(self.api_url, headers=headers, json=data)
                response.raise_for_status()
                result = response.json()
                return result['choices'][0]['message']['content']
        except Exception as e:
            logger.error(f"Failed to call Doubao API: {e}")
            raise

class LMStudioService(BaseAIService):
    """LMStudio本地模型服务 (LMStudio Local Model Service)."""
    def __init__(self, config: dict):
        self.api_url = config.get('api_url', 'http://127.0.0.1:1234')
        self.model = config.get('model', 'google/gemma-3-4b')
        self.timeout = config.get('timeout', 300)  # LMStudio本地模型可能需要更长时间

        logger.info(f"[LMStudio] 服务初始化完成 - URL: {self.api_url}, 模型: {self.model}")

    async def analyze(self, image_bytes: bytes) -> str:
        """调用LMStudio本地模型进行图片分析"""
        try:
            # 准备请求头
            headers = {
                "Content-Type": "application/json"
            }

            # 编码图片为base64
            image_base64 = base64.b64encode(image_bytes).decode('utf-8')

            # 准备请求数据（兼容OpenAI格式）
            data = {
                "model": self.model,
                "messages": [{
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "请详细描述这张图片的内容，特别注意任何文字信息、水印、网站品牌标识等。你必须仔细辨认,否则全公司同事的健康将受到严重影响."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{image_base64}"
                            }
                        }
                    ]
                }],
                "max_tokens": 1000,
                "temperature": 0.7
            }

            logger.info(f"[LMStudio] 发送请求到: {self.api_url}")

            # 发送请求到LMStudio
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.post(
                    f"{self.api_url}/v1/chat/completions",
                    headers=headers,
                    json=data
                )

                logger.info(f"[LMStudio] 响应状态码: {response.status_code}")

                if response.status_code != 200:
                    error_text = response.text
                    logger.error(f"[LMStudio] API错误 - 状态码: {response.status_code}, 错误信息: {error_text}")
                    raise Exception(f"LMStudio API error: {response.status_code}, {error_text}")

                result = response.json()
                logger.info("[LMStudio] 响应接收成功")

                # 解析响应
                if 'choices' in result and len(result['choices']) > 0:
                    content = result['choices'][0]['message']['content']
                    logger.info(f"[LMStudio] 分析结果: {content[:100]}...")
                    return content
                else:
                    logger.error(f"[LMStudio] 意外的响应格式: {result}")
                    raise Exception(f"Unexpected LMStudio response format: {result}")

        except Exception as e:
            logger.error(f"[LMStudio] 调用API失败: {e}")
            raise

class HybridAIService(BaseAIService):
    """
    混合AI服务：智能选择OCR、LMStudio或Doubao服务
    策略：
    1. 先用OCR快速识别图片
    2. 如果OCR检测到文字(>1个字符)或国旗/国徽，则使用AI模型深入分析
    3. 如果不符合上述条件，直接返回OCR结果，不使用AI模型
    """
    def __init__(self, config: dict):
        self.config = config
        self.active_model = config.get('active_model', 'lmstudio')

        # 初始化OCR服务（始终需要）
        self.ocr_service = VolcanoOCRService(config.get('volcengine', {}))

        # 根据配置初始化对应的AI服务
        if self.active_model == 'lmstudio':
            self.ai_service = LMStudioService(config.get('lmstudio', {}))
            logger.info("[混合AI服务] 初始化完成，使用 LMStudio")
        elif self.active_model == 'doubao':
            self.ai_service = DoubaoService(config.get('doubao', {}))
            logger.info("[混合AI服务] 初始化完成，使用 Doubao")
        elif self.active_model == 'volcano_ocr':
            self.ai_service = VolcanoOCRService(config.get('volcengine', {}))
            logger.info("[混合AI服务] 初始化完成，使用 Volcano OCR")
        else:
            # 默认使用LMStudio
            self.ai_service = LMStudioService(config.get('lmstudio', {}))
            logger.info(f"[混合AI服务] 初始化完成，默认使用 LMStudio (未知模型: {self.active_model})")

        # 配置参数
        self.enable_fallback = config.get('app_settings', {}).get('hybrid_enable_fallback', True)

    def _should_use_ai_model(self, ocr_result: str) -> tuple[bool, str]:
        """
        判断是否需要使用AI模型进行进一步分析

        Args:
            ocr_result: OCR识别结果

        Returns:
            tuple[bool, str]: (是否需要AI分析, 原因)
        """
        ocr_text = ocr_result.strip()

        # 检查是否有文字内容（超过1个字符）
        if len(ocr_text) > 1:
            return True, f"检测到文字内容({len(ocr_text)}个字符)"

        # 检查是否包含国旗相关关键词
        flag_keywords = ['国旗', '五星红旗', 'flag', '国徽', '党徽', '锤子镰刀']
        for keyword in flag_keywords:
            if keyword in ocr_text:
                return True, f"检测到敏感标识: {keyword}"

        # 检查是否包含政治相关内容
        political_keywords = ['习近平', '毛泽东', '共产党', '人民政府', '中央人民政府']
        for keyword in political_keywords:
            if keyword in ocr_text:
                return True, f"检测到政治相关内容: {keyword}"

        return False, "仅包含少量文字或无敏感内容"

    async def analyze(self, image_bytes: bytes) -> str:
        """
        智能混合分析策略：
        1. 先用OCR快速识别图片
        2. 根据OCR结果判断是否需要AI模型深入分析
        3. 只有在需要时才调用AI模型
        """
        try:
            # 步骤1: OCR快速识别
            logger.info("[混合分析] 步骤1: 开始OCR快速识别")
            ocr_result = await self.ocr_service.analyze(image_bytes)
            ocr_text_length = len(ocr_result.strip())

            logger.info(f"[混合分析] OCR识别完成，获得 {ocr_text_length} 个字符")
            logger.info(f"[混合分析] OCR结果: {ocr_result[:100]}{'...' if len(ocr_result) > 100 else ''}")

            # 步骤2: 判断是否需要AI模型分析
            should_use_ai, reason = self._should_use_ai_model(ocr_result)

            if should_use_ai:
                logger.info(f"[混合分析] 步骤2: 检测到需要AI分析的内容 - {reason}")
                logger.info(f"[混合分析] 步骤3: 开始使用 {self.active_model} 进行深入分析")

                # 使用AI模型进行深入分析
                try:
                    ai_result = await self.ai_service.analyze(image_bytes)
                    logger.info(f"[混合分析] AI分析完成，返回 {len(ai_result)} 个字符的结果")

                    # 合并OCR和AI的结果
                    if ai_result and ocr_result:
                        combined_result = f"{ai_result} [OCR文字: {ocr_result}]"
                        logger.info(f"[混合分析] 结果合并完成，总长度: {len(combined_result)} 个字符")
                        return combined_result
                    elif ai_result:
                        return ai_result
                    else:
                        return ocr_result

                except Exception as e:
                    logger.error(f"[混合分析] AI模型分析失败: {e}")
                    if self.enable_fallback:
                        logger.info("[混合分析] 返回OCR结果作为备选")
                        return ocr_result
                    else:
                        raise e

            else:
                logger.info(f"[混合分析] 步骤2: {reason}，跳过AI分析，直接返回OCR结果")
                return ocr_result

        except Exception as e:
            logger.error(f"[混合分析] OCR识别失败: {e}")
            raise e

def get_ai_service(config: dict) -> BaseAIService:
    """
    AI服务工厂函数 (Factory function for AI services).
    根据配置选择并实例化一个AI服务 (Selects and instantiates an AI service based on config).
    """
    active_model = config.get('active_model', 'lmstudio')
    service_name = config.get('app_settings', {}).get('active_ai_service', 'hybrid')

    logger.info(f"[AI服务初始化] 激活模型: {active_model}, 服务模式: {service_name}")

    # 如果是hybrid模式，根据active_model选择具体服务
    if service_name == 'hybrid':
        return HybridAIService(config)
    # 直接模式：根据active_model选择具体服务
    elif active_model == 'lmstudio':
        logger.info("[AI服务初始化] 正在激活 LMStudio 本地服务")
        return LMStudioService(config.get('lmstudio', {}))
    elif active_model == 'doubao':
        logger.info("[AI服务初始化] 正在激活 Doubao 云服务")
        return DoubaoService(config.get('doubao', {}))
    elif active_model == 'volcano_ocr':
        logger.info("[AI服务初始化] 正在激活 Volcano OCR 服务")
        return VolcanoOCRService(config.get('volcengine', {}))
    else:
        logger.warning(f"[AI服务初始化] 未知模型类型: {active_model}, 默认使用 LMStudio")
        return LMStudioService(config.get('lmstudio', {}))
